#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_GRAD, typename TYPE_X, typename TYPE_TANHX, typename TYPE_Y>
class MishGradKernel {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> gradQueue;
    TQue<QuePosition::VECIN, BUFFER_NUM> xQueue;
    TQue<QuePosition::VECIN, BUFFER_NUM> tanhxQueue;
    TQue<QuePosition::VECOUT, BUFFER_NUM> yQueue;

    GlobalTensor<TYPE_X> gradGm;
    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_X> tanhxGm;
    GlobalTensor<TYPE_Y> yGm;
    TBuf<QuePosition::VECCALC> Cal1, Cal2, Tmp1, Tmp2;
    // TBuf<QuePosition::VECCALC> Tmp_F, Row_F, Col_F, Lidx_F, Gidx_F;

    uint32_t ALIGN_NUM;
    uint32_t totalLength;
    uint32_t tileLength;
    uint32_t tileNum;
    uint32_t lasttileLength;

public:
    __aicore__ inline MishGradKernel() {}
    __aicore__ inline void Init(GM_ADDR grad, GM_ADDR x, GM_ADDR tanhx, GM_ADDR y, uint32_t ALIGN_NUM, uint32_t totalLength,
                                uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength)

    {
        this->ALIGN_NUM = ALIGN_NUM;
        this->totalLength = totalLength;
        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;

        gradGm.SetGlobalBuffer((__gm__ TYPE_X *)grad, totalLength);
        xGm.SetGlobalBuffer((__gm__ TYPE_X *)x, totalLength);
        tanhxGm.SetGlobalBuffer((__gm__ TYPE_X *)tanhx, totalLength);
        yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y, totalLength);

        pipe.InitBuffer(gradQueue, BUFFER_NUM, tileLength * sizeof(TYPE_X)); //自动向上32Byte对齐
        pipe.InitBuffer(xQueue, BUFFER_NUM, tileLength * sizeof(TYPE_X)); //自动向上32Byte对齐
        pipe.InitBuffer(tanhxQueue, BUFFER_NUM, tileLength * sizeof(TYPE_X)); //自动向上32Byte对齐
        pipe.InitBuffer(yQueue, BUFFER_NUM, tileLength * sizeof(TYPE_Y));

        pipe.InitBuffer(Cal1, tileLength * sizeof(TYPE_X));
        pipe.InitBuffer(Cal2, tileLength * sizeof(TYPE_X));
        // pipe.InitBuffer(Cal3, tileLength * sizeof(TYPE_X));
        pipe.InitBuffer(Tmp1, tileLength * sizeof(TYPE_X));
        pipe.InitBuffer(Tmp2, tileLength * sizeof(TYPE_X));
        // pipe.InitBuffer(Tmp3, tileLength * sizeof(TYPE_X));

    }

    __aicore__ inline void Process()
    {
        uint32_t  loopCount = tileNum;
        
        for (int32_t i = 0; i < loopCount - 1; i++) 
        {                
            CopyIn(i, tileLength);
            Compute(i, tileLength);
            CopyOut(i, tileLength);
        }

        CopyIn(tileNum - 1, lasttileLength);
        Compute(tileNum - 1, lasttileLength);
        CopyOut(tileNum - 1, lasttileLength);
            
    }
private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> gradLocal = gradQueue.AllocTensor<TYPE_X>();
        LocalTensor<TYPE_X> xLocal = xQueue.AllocTensor<TYPE_X>();
        LocalTensor<TYPE_X> tanhxLocal = tanhxQueue.AllocTensor<TYPE_X>();
        uint32_t startindex = progress * tileLength;
        DataCopy(gradLocal, gradGm[startindex], length);
        DataCopy(xLocal, xGm[startindex], length);
        DataCopy(tanhxLocal, tanhxGm[startindex], length);
        gradQueue.EnQue(gradLocal);
        xQueue.EnQue(xLocal);
        tanhxQueue.EnQue(tanhxLocal);
    }
    /*
        if input_tanhx is None:
            exp_x = np.exp(input_x)
            exp_add_x = exp_x + 1
            rec_exp_add_x = 1.0 / (exp_add_x * exp_add_x + 1.0)
            result_1 = exp_add_x * exp_x * input_x * rec_exp_add_x * rec_exp_add_x * 4
            result = result_1 - 2 * rec_exp_add_x + 1
        else:
            pow_input_tanhx = input_tanhx * input_tanhx
            result = input_tanhx + input_x * (1 - pow_input_tanhx) * np.exp(input_x) / (1 + np.exp(input_x))
        result = result * input_grad
        return [result, ]

     */

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> gradLocal = gradQueue.DeQue<TYPE_X>();
        LocalTensor<TYPE_X> xLocal = xQueue.DeQue<TYPE_X>();
        LocalTensor<TYPE_X> tanhxLocal = tanhxQueue.DeQue<TYPE_X>();
        LocalTensor<TYPE_Y> yLocal = yQueue.AllocTensor<TYPE_Y>();

        LocalTensor<TYPE_X> cal1 = Cal1.Get<TYPE_X>();
        LocalTensor<TYPE_X> cal2 = Cal2.Get<TYPE_X>();
        // LocalTensor<TYPE_X> cal3 = Cal3.Get<TYPE_X>();
        LocalTensor<TYPE_X> tmp1 = Tmp1.Get<TYPE_X>();
        LocalTensor<TYPE_X> tmp2 = Tmp2.Get<TYPE_X>();
        // LocalTensor<TYPE_X> tmp3 = Tmp3.Get<TYPE_X>();

        Mul(cal1, tanhxLocal, tanhxLocal, length);  //tanx^2
        Exp(tmp1, xLocal, length);  //e^x
        Muls(cal1, cal1, static_cast<TYPE_X>(-1.0), length); // -tanx^2
        Adds(cal1, cal1, static_cast<TYPE_X>(1.0), length); //1 - tanx^2
        Mul(cal1, cal1, tmp1, length);  //(1 - pow_input_tanhx) * np.exp(input_x)
        Mul(cal1, cal1, xLocal, length);  // input_x * (1 - pow_input_tanhx) * np.exp(input_x)
        
        Adds(tmp1, tmp1, static_cast<TYPE_X>(1.0), length); //(1 + np.exp(input_x))
        Div(cal1, cal1, tmp1, length);  //input_x * (1 - pow_input_tanhx) * np.exp(input_x) / (1 + np.exp(input_x))
        Add(cal1, cal1, tanhxLocal, length);    //input_tanhx + input_x * (1 - pow_input_tanhx) * np.exp(input_x) / (1 + np.exp(input_x))
        Mul(yLocal, cal1, gradLocal, length);

        gradQueue.FreeTensor(gradLocal);
        xQueue.FreeTensor(xLocal);
        tanhxQueue.FreeTensor(tanhxLocal);
        yQueue.EnQue<DTYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, uint32_t length)
    {
        LocalTensor<DTYPE_Y> yLocal = yQueue.DeQue<DTYPE_Y>();
        uint32_t startindex = progress * tileLength;
        DataCopy(yGm[startindex], yLocal, length);
        yQueue.FreeTensor(yLocal);
    }

};

extern "C" __global__ __aicore__ void mish_grad(GM_ADDR grad, GM_ADDR x, GM_ADDR tanhx, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    // if(tiling_data.optional == 3)
    // {
        MishGradKernel<DTYPE_GRAD, DTYPE_X, DTYPE_TANHX, DTYPE_Y> op;
        op.Init(grad, x, tanhx, y, tiling_data.ALIGN_NUM, tiling_data.totalLength,
            tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileLength);
        op.Process();
    // }
}